package task;


import cn.xueden.common.core.context.ContextHolder;
import cn.xueden.common.core.context.RequestContext;
import com.sun.istack.internal.NotNull;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * @author Gyg
 * @date 2022/1/20 10:47
 */
@Slf4j
public class ContextTaskDecorator implements TaskDecorator {

    @Override
    @NotNull
    public Runnable decorate(@NotNull Runnable runnable) {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        if (requestAttributes == null) {
            return () -> {
                log.debug("开启子线程的父线程本身就没有上下文！！！");
                runnable.run();
            };
        }
        ServletRequestAttributes context = (ServletRequestAttributes)RequestContextHolder.currentRequestAttributes();
        Map<String, String> headMap = getHeadMap(context);
        RequestContext requestContext = ContextHolder.get();
        return () -> {
            try {
                RequestContextHolder.setRequestAttributes(context, true);
                HeaderContextHolder.setContext(headMap);
                ContextHolder.set(requestContext);
                runnable.run();
            } finally {
                RequestContextHolder.resetRequestAttributes();
                HeaderContextHolder.removeContext();
                ContextHolder.remove();
            }
        };
    }

    private Map<String, String> getHeadMap(ServletRequestAttributes context) {
        Map<String, String> headMap = new HashMap<>();
        HttpServletRequest request = context.getRequest();
        Enumeration<String> headerNames = request.getHeaderNames();
        if (headerNames != null) {
            while (headerNames.hasMoreElements()) {
                String key = headerNames.nextElement();
                String value = request.getHeader(key);
                headMap.put(key, value);
            }
        }
        return headMap;
    }

}
